package netfilter

import (
	"fmt"
	lp "lbs-sdk/protocol"
	utils "lbs-sdk/utils"
	"sync"
	"sync/atomic"
	"time"

	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/header"
)

// IP会话信息
type IpSession struct {
	RealSrcIP   tcpip.Address                 // 真实源IP
	RealSrcPort uint16                        // 真实源端口
	FakeIP      tcpip.Address                 // 固定的fakeIP (10.200.1.121)
	FakePort    uint16                        // 分配的fake端口
	DstIP       tcpip.Address                 // 目标IP
	DstPort     uint16                        // 目标端口
	Protocol    tcpip.TransportProtocolNumber // 协议类型
	LastSeen    time.Time                     // 最后活动时间 (Unix时间戳)
	Options     NfIpPacketOptions
}

// 简化的NAT表 - 只需要一个65536的数组
type FastNatTable struct {
	tcpTable  [65536]*IpSession // TCP协议的端口映射
	udpTable  [65536]*IpSession // UDP协议的端口映射
	icmpTable [65536]*IpSession // ICMP协议的端口映射
	mutex     sync.RWMutex
}

var DefaultNatTableIpv4 = NewFastNatTable()
var DefaultNatTableIpv6 = NewFastNatTable()

func NewFastNatTable() *FastNatTable {
	return &FastNatTable{}
}

func GetNatTable(ipVersion uint8) *FastNatTable {
	switch ipVersion {
	case header.IPv4Version:
		return DefaultNatTableIpv4
	case header.IPv6Version:
		return DefaultNatTableIpv6
	default:
		panic("unknown ip version")
	}
}

// 根据协议获取对应的表
func (nt *FastNatTable) getTableByProtocol(protocol tcpip.TransportProtocolNumber) *[65536]*IpSession {
	switch protocol {
	case header.TCPProtocolNumber: // TCP
		return &nt.tcpTable
	case header.UDPProtocolNumber: // UDP
		return &nt.udpTable
	case header.ICMPv4ProtocolNumber, header.ICMPv6ProtocolNumber: // ICMP
		return &nt.icmpTable
	default:
		return nil
	}
}

// 查找可用端口 (从指定端口开始往后找)
func (nt *FastNatTable) findAvailableIdx(table *[65536]*IpSession, realSrcIP tcpip.Address, startPort uint16) (uint16, bool) {
	// 从startPort开始查找
	for port := startPort; port != 0; port++ {
		if table[port] == nil {
			return port, false
		} else {
			if table[port].RealSrcIP == realSrcIP {
				return port, true
			}
		}
	}

	// 如果从startPort到65535都被占用，从1开始查找
	for port := uint16(1); port < startPort; port++ {
		if table[port] == nil {
			return port, false
		}
	}

	return 0, false // 所有端口都被占用
}

// 添加NAT条目
func (nt *FastNatTable) AddSession(options NfIpPacketOptions, fakeIp tcpip.Address, realSrcIP tcpip.Address, realSrcPort uint16, dstIP tcpip.Address, dstPort uint16, protocol tcpip.TransportProtocolNumber) (*IpSession, bool) {
	nt.mutex.Lock()
	defer nt.mutex.Unlock()

	// 1. 根据协议获取对应的表
	table := nt.getTableByProtocol(protocol)

	// 从startPort开始查找
	idx, reuse := nt.findAvailableIdx(table, realSrcIP, realSrcPort)
	if idx == 0 {
		return nil, false // 没有可用端口
	}
	if reuse {
		table[idx].LastSeen = time.Now()
		// fmt.Printf("^ %d add sess reuse id:%d tb:%p nt:%p\n", protocol, idx, table, nt)
		return table[idx], true // 复用已有会话
	}

	// 3. 创建会话
	session := &IpSession{
		RealSrcIP:   realSrcIP,
		RealSrcPort: realSrcPort,
		FakeIP:      fakeIp,
		FakePort:    uint16(idx),
		DstIP:       dstIP,
		DstPort:     dstPort,
		Protocol:    protocol,
		LastSeen:    time.Now(),
		Options:     options,
	}
	// 4. 存储到表中
	table[idx] = session
	return session, true
}

// 根据fakePort查找会话 (收到返回包时使用)
func (nt *FastNatTable) GetByFakePort(fakePort uint16, protocol tcpip.TransportProtocolNumber) (*IpSession, bool) {
	nt.mutex.RLock()
	defer nt.mutex.RUnlock()

	table := nt.getTableByProtocol(protocol)
	session := table[fakePort]
	fmt.Printf("v %d get session by fakePort %d tb:%p nt:%p\n", protocol, fakePort, table, nt)
	// 更新最后活动时间
	if session != nil {
		session.LastSeen = time.Now()
	}
	return session, session != nil
}

type NatTable struct {
}

var Ipv4Snmp utils.IpPkgSnmp

func PerformDnatIpv4(ipv4Hdr header.IPv4, oldPPort uint16) (action lp.NETFILTER_ACTION, session *IpSession) {
	// 初始化返回值
	action = lp.NETFILTER_DROP

	var transportProtocol tcpip.TransportProtocolNumber
	var ipPayload []byte
	var oldDst tcpip.Address
	var ok bool
	if len(ipv4Hdr) < header.IPv4MinimumSize {
		atomic.AddUint64(&Ipv4Snmp.HeaderTooShort, 1)
		return
	}
	transportProtocol = ipv4Hdr.TransportProtocol()
	oldDst = ipv4Hdr.DestinationAddress()
	ipPayload = ipv4Hdr.Payload()
	tb := GetNatTable(header.IPv4Version)
	if tb == nil {
		atomic.AddUint64(&Ipv4Snmp.SenderNil, 1)
		return
	}
	// 传输层处理
	switch transportProtocol {
	case header.TCPProtocolNumber:
		if len(ipPayload) < header.TCPMinimumSize {
			atomic.AddUint64(&Ipv4Snmp.TcpHeaderTooShort, 1)
			return
		}
		tcp := header.TCP(ipPayload)
		oldDstPort := tcp.DestinationPort()
		if oldPPort > 0 {
			oldDstPort = oldPPort
		}
		session, ok = tb.GetByFakePort(oldDstPort, header.TCPProtocolNumber)
		if !ok {
			atomic.AddUint64(&Ipv4Snmp.NoSessionFound, 1)
			return
		}
		newDstPort := session.RealSrcPort
		newDst := session.RealSrcIP
		ipv4Hdr.SetDestinationAddressWithChecksumUpdate(newDst)
		tcp.SetDestinationPortWithChecksumUpdate(newDstPort)
		tcp.UpdateChecksumPseudoHeaderAddress(oldDst, newDst, true)
	case header.UDPProtocolNumber:
		if len(ipPayload) < header.UDPMinimumSize {
			atomic.AddUint64(&Ipv4Snmp.UdpHeaderTooShort, 1)
			return
		}
		udp := header.UDP(ipPayload)
		oldDstPort := udp.DestinationPort()
		if oldPPort > 0 {
			oldDstPort = oldPPort
		}
		session, ok = tb.GetByFakePort(oldDstPort, header.UDPProtocolNumber)
		if !ok {
			fmt.Printf("v udp no session found for port %d\n", oldDstPort)
			atomic.AddUint64(&Ipv4Snmp.NoSessionFound, 1)
			return
		}
		newDstPort := session.RealSrcPort
		newDst := session.RealSrcIP
		ipv4Hdr.SetDestinationAddressWithChecksumUpdate(newDst)
		udp.SetDestinationPortWithChecksumUpdate(newDstPort)
		udp.UpdateChecksumPseudoHeaderAddress(oldDst, newDst, true)
	case header.ICMPv4ProtocolNumber:
		if len(ipv4Hdr.Payload()) < header.ICMPv4MinimumSize {
			atomic.AddUint64(&Ipv4Snmp.IcmpHeaderTooShort, 1)
			return
		}
		icmp := header.ICMPv4(ipv4Hdr.Payload())
		if icmp.Type() == header.ICMPv4Echo {
			return
		}
		oldIdent := icmp.Ident()
		if oldPPort > 0 {
			oldIdent = oldPPort
		}
		session, ok = tb.GetByFakePort(oldIdent, header.ICMPv4ProtocolNumber)
		if !ok {
			atomic.AddUint64(&Ipv4Snmp.NoSessionFound, 1)
			return
		}

		ipv4Hdr.SetDestinationAddressWithChecksumUpdate(session.RealSrcIP)
		icmp.SetIdentWithChecksumUpdate(session.RealSrcPort)
	default:
		atomic.AddUint64(&Ipv4Snmp.NextHeaderNotSupport, 1)
		return
	}
	action = lp.NETFILTER_ACCEPT
	return
}

func ProcessIpv4(ipv4Hdr header.IPv4, oldPort uint16) (session *IpSession) {
	_, session = PerformDnatIpv4(ipv4Hdr, oldPort)
	if session == nil {
		return
	}
	return
}
